package example.NIOTCPTest;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class FixedNioServer {
    private Selector selector;
    private ServerSocketChannel serverChannel;
    private Charset charset = Charset.forName("UTF-8");

    // 为每个客户端连接维护一个会话状态
    private Map<SocketChannel, ClientSession> clientSessions = new HashMap<>();

    class ClientSession {
        ByteBuffer readBuffer = ByteBuffer.allocate(5); // 读取缓冲区
        ByteBuffer lengthBuffer = ByteBuffer.allocate(4);  // 专门用于读取长度头
        boolean readingLength = true; // 当前是否在读取长度头
        int messageLength = 0;        // 当前消息的长度
        ByteBuffer messageBuffer;     // 用于存储消息体

        void reset() {
            readingLength = true;
            messageLength = 0;
            lengthBuffer.clear();
            messageBuffer = null;
        }
    }

    public void start(int port) throws IOException {
        selector = Selector.open();
        serverChannel = ServerSocketChannel.open();
        serverChannel.configureBlocking(false);
        serverChannel.socket().bind(new InetSocketAddress(port));
        serverChannel.register(selector, SelectionKey.OP_ACCEPT);

        System.out.println("NIO 服务器启动，监听端口 " + port);

        while (true) {
            selector.select();
            Iterator<SelectionKey> keyIterator = selector.selectedKeys().iterator();

            while (keyIterator.hasNext()) {
                SelectionKey key = keyIterator.next();
                keyIterator.remove();

                try {
                    if (key.isAcceptable()) {
                        handleAccept(key);
                    } else if (key.isReadable()) {
                        handleRead(key);
                    }
                } catch (IOException e) {
                    System.err.println("处理连接时出错: " + e.getMessage());
                    handleDisconnect(key);
                }
            }
        }
    }

    private void handleAccept(SelectionKey key) throws IOException {
        ServerSocketChannel serverChannel = (ServerSocketChannel) key.channel();
        SocketChannel clientChannel = serverChannel.accept();
        clientChannel.configureBlocking(false);
        clientChannel.register(selector, SelectionKey.OP_READ);

        // 为新客户端创建会话
        clientSessions.put(clientChannel, new ClientSession());
        System.out.println("客户端连接: " + clientChannel.getRemoteAddress());
    }

    private void handleRead(SelectionKey key) throws IOException {
        SocketChannel clientChannel = (SocketChannel) key.channel();
         ClientSession session = clientSessions.get(clientChannel);

        if (session == null) {
            return;
        }

        // 读取数据到缓冲区
        int bytesRead = clientChannel.read(session.readBuffer);

        if (bytesRead == -1) {
            handleDisconnect(key);
            return;
        }

        if (bytesRead > 0) {
            // 处理接收到的数据
            processReceivedData(clientChannel, session);
        }
    }

    private void processReceivedData(SocketChannel clientChannel, ClientSession session) {
        // 准备读取数据
        session.readBuffer.flip();

        try {
            while (session.readBuffer.hasRemaining()) {
                if (session.readingLength) {
                    // 读取长度头
                    int remaining = Math.min(session.readBuffer.remaining(), session.lengthBuffer.remaining());
                    for (int i = 0; i < remaining; i++) {
                        session.lengthBuffer.put(session.readBuffer.get());
                    }

                    // 检查是否已读取完整的4字节长度头
                    if (!session.lengthBuffer.hasRemaining()) {
                        session.lengthBuffer.flip();
                        session.messageLength = session.lengthBuffer.getInt();
                        System.out.println("期望消息长度: " + session.messageLength);

                        if (session.messageLength <= 0 || session.messageLength > 1024 * 1024) {
                            System.err.println("无效的消息长度: " + session.messageLength);
                            session.reset();
                            continue;
                        }

                        // 创建消息缓冲区
                        session.messageBuffer = ByteBuffer.allocate(session.messageLength);
                        session.readingLength = false;
                        session.lengthBuffer.clear();
                    }
                } else {
                    // 读取消息体
                    int remaining = Math.min(session.readBuffer.remaining(),
                            session.messageBuffer.remaining());
                    for (int i = 0; i < remaining; i++) {
                        session.messageBuffer.put(session.readBuffer.get());
                    }

                    // 检查是否已读取完整的消息体
                    if (!session.messageBuffer.hasRemaining()) {
                        // 完整消息已接收
                        session.messageBuffer.flip();
                        String fullMessage = charset.decode(session.messageBuffer).toString();
                        System.out.println("收到完整消息(" + fullMessage.length() + "字符): " + fullMessage);

                        // 发送响应
                        sendResponse(clientChannel, "已收到: " + fullMessage.length() + " 字符");

                        // 重置会话状态，准备接收下一条消息
                        session.reset();
                    }
                }
            }
        } finally {
            // 压缩缓冲区，保留未处理的数据
            session.readBuffer.compact();
        }
    }

    private void sendResponse(SocketChannel clientChannel, String response) {
        try {
            byte[] responseData = response.getBytes(charset);
            ByteBuffer responseBuffer = ByteBuffer.allocate(4 + responseData.length);

            // 写入长度头
            responseBuffer.putInt(responseData.length);
            // 写入消息体
            responseBuffer.put(responseData);
            responseBuffer.flip();

            while (responseBuffer.hasRemaining()) {
                clientChannel.write(responseBuffer);
            }

            System.out.println("发送响应: " + response);
        } catch (IOException e) {
            System.err.println("发送响应失败: " + e.getMessage());
        }
    }

    private void handleDisconnect(SelectionKey key) {
        SocketChannel clientChannel = (SocketChannel) key.channel();
        clientSessions.remove(clientChannel);

        try {
            System.out.println("客户端断开连接: " + clientChannel.getRemoteAddress());
            clientChannel.close();
        } catch (IOException e) {
            // 忽略关闭异常
        }
        key.cancel();
    }

    public static void main(String[] args) {
        try {
            new FixedNioServer().start(8888);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}